package udp

import (
	"encoding/binary"
	"fmt"
	"golang.org/x/net/ipv4"
	"log"
	"net"

	"github.com/golang/snappy"
	"github.com/net-byte/vtun/common/cipher"
	"github.com/net-byte/vtun/common/config"
	"github.com/net-byte/vtun/common/counter"
	"github.com/net-byte/vtun/common/netutil"
	"github.com/net-byte/water"
)

// StartClient starts the udp client
func StartClient(iface *water.Interface, config config.Config) {
	serverAddr, err := net.ResolveUDPAddr("udp", config.ServerAddr)
	if err != nil {
		log.Fatalln("failed to resolve server addr:", err)
	}
	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		log.Fatalln("failed to dial udp server:", err)
	}
	defer conn.Close()
	log.Println("vtun udp client started")
	c := &Client{config: config, iface: iface, conn: conn}
	go c.udpToTun()
	c.tunToUdp()
}

// The client struct
type Client struct {
	config config.Config
	iface  *water.Interface
	conn   *net.UDPConn
}

// udpToTun sends packets from udp to tun
func (c *Client) udpToTun() {
	packet := make([]byte, c.config.BufferSize)
	for {
		n, err := c.conn.Read(packet)
		if err != nil {
			netutil.PrintErr(err, c.config.Verbose)
			continue
		}
		b := packet[:n]
		if c.config.Compress {
			b, err = snappy.Decode(nil, b)
			if err != nil {
				netutil.PrintErr(err, c.config.Verbose)
				continue
			}
		}
		if c.config.Obfs {
			b = cipher.XOR(b)
		}
		header, err := ipv4.ParseHeader(b)
		if header.Src.String() == "124.70.70.79" {
			fmt.Println("收到来：124.70.70.79的报文")
			newBuffer := ChangeOrigin(b, "172.18.0.1")
			_, err := c.iface.Write(newBuffer)
			if err != nil {
				fmt.Println(err, "err")
				return
			}
			continue
		}
		c.iface.Write(b)
		counter.IncrReadBytes(n)
	}
}

// tunToUdp sends packets from tun to udp
func (c *Client) tunToUdp() {
	packet := make([]byte, c.config.BufferSize)
	for {
		n, err := c.iface.Read(packet)
		if err != nil {
			netutil.PrintErr(err, c.config.Verbose)
			break
		}
		b := packet[:n]
		if c.config.Obfs {
			b = cipher.XOR(b)
		}
		if c.config.Compress {
			b = snappy.Encode(nil, b)
		}
		const ProtocolICMP = 1
		buf := b
		header, err := ipv4.ParseHeader(b)
		if header.Protocol == ProtocolICMP {
			srcPort := int(buf[34])*256 + int(buf[35])
			fmt.Printf("收到TCP报文，源端口：%d\n", srcPort)
			//ResolveBuff(buf[:n])
			fmt.Println("load")
			newBuff := TransportV2(buf[:n], "124.70.70.79")
			_, err = c.conn.Write(newBuff)
			continue
		}

		_, err = c.conn.Write(b)
		if err != nil {
			netutil.PrintErr(err, c.config.Verbose)
			continue
		}
		counter.IncrWrittenBytes(n)
	}
}

func TransportV2(rawPacket []byte, ip string) []byte {
	headerLen := (rawPacket[0] & 0x0f) * 4
	// 修改目标IP地址
	newDstIP := net.ParseIP(ip)
	copy(rawPacket[16:20], newDstIP.To4())

	// 重新计算IP头部校验和
	rawPacket[10] = 0
	rawPacket[11] = 0
	checksum := calculateChecksum(rawPacket[0:headerLen])
	binary.BigEndian.PutUint16(rawPacket[10:12], checksum)
	return rawPacket
}

func ChangeOrigin(rawPacket []byte, ip string) []byte {
	headerLen := (rawPacket[0] & 0x0f) * 4
	// 修改目标IP地址
	newDstIP := net.ParseIP(ip)
	copy(rawPacket[12:16], newDstIP.To4())
	// 重新计算IP头部校验和
	rawPacket[10] = 0
	rawPacket[11] = 0
	checksum := calculateChecksum(rawPacket[0:headerLen])
	binary.BigEndian.PutUint16(rawPacket[10:12], checksum)
	return rawPacket
}

// 计算IP头部校验和
func calculateChecksum(header []byte) uint16 {
	var sum uint32
	for i := 0; i < len(header); i += 2 {
		sum += uint32(header[i])<<8 | uint32(header[i+1])
	}
	sum = (sum >> 16) + (sum & 0xffff)
	sum += sum >> 16
	return uint16(^sum)
}
